# import json 
# def isCorrectResult(input_str,result):
#     # 1.确定result是json格式
#     try:
#         result_json=json.loads(result)
#     except Exception as e:
#         print('result不能被解析为json格式',e)
#         return False

#     # 2.确定result的键是齐全的
#     check_list=['id','question','answer']
#     is_keys_all_in_result = all(item in result_json for item in check_list)
#     if not is_keys_all_in_result:
#         print('result内容缺失')
#         return False
    
#     # 3.保证result的id和原始id一致（简单的内容检查）
#     is_id_same=json.loads(input_str)['id']==result_json['id']
#     if not is_id_same:
#         print('result的id错误')
#         return False
    
#     return result

# # test case
# # 1. correct
# print(isCorrectResult('{"id": 8, "question": "我想申请教学改革研究项目。", "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}','{"id": 8, "question": "要求啥条件？", "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}'))
# # 2. not json
# print(isCorrectResult('{"id": 8, "question": "我想申请教学改革研究项目。", "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}','"id": 8, "question": "要求啥条件？", "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}'))
# # 3. key missing
# print(isCorrectResult('{"id": 8, "question": "我想申请教学改革研究项目。", "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}','{"id": 8,  "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}'))
# # 4. id mismatch
# print(isCorrectResult('{"id": 9, "question": "我想申请教学改革研究项目。", "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}','{"id": 8,  "answer": "申请本科教学改革研究项目请参考这个网址https://ehall.neu.edu.cn//db_portal/guide?id=2EC05B47-1355-4EC6-86DB-B81698081B1D"}'))

from dataclasses import dataclass
from typing import Any
from transformers import BertTokenizer, BertForNextSentencePrediction,Trainer,TrainingArguments,DataCollatorWithPadding
from torch.utils.data import Dataset,DataLoader
from functools import partial
import json
import torch
tokenizer = BertTokenizer.from_pretrained("BAAI/bge-large-zh-v1.5")
model=BertForNextSentencePrediction.from_pretrained("BAAI/bge-large-zh-v1.5")

tokenizer.save_pretrained('/home/lxy/multiR/NSPtraining/bge')
model.save_pretrained('/home/lxy/multiR/NSPtraining/bge')